package com.software.usercenter.judge;

import java.lang.reflect.Field;
import java.security.Permission;

public class MemorySecurityManager extends SecurityManager {

    private final long maxMemory;

    public MemorySecurityManager(long maxMemory) {
        this.maxMemory = maxMemory;
    }

    @Override
    public void checkExec(String cmd) {
        // 限制程序的内存使用
        limitMemoryUsage();

        super.checkExec(cmd);
    }

    private void limitMemoryUsage() {
        long maxMemoryBytes = maxMemory;
        try {
            // 获取 Java 虚拟机内部的 Unsafe 对象
            Field field = sun.misc.Unsafe.class.getDeclaredField("theUnsafe");
            field.setAccessible(true);
            sun.misc.Unsafe unsafe = (sun.misc.Unsafe) field.get(null);

            // 分配一块内存，并将其全部初始化为 0
            long memoryAddress = unsafe.allocateMemory(maxMemoryBytes);
            unsafe.setMemory(memoryAddress, maxMemoryBytes, (byte) 0);

            // 将限制内存的方法注册到虚拟机关闭钩子中，在程序退出时释放内存
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                unsafe.freeMemory(memoryAddress);
            }));
        } catch (Exception e) {
            // 改成自定义异常
            throw new RuntimeException(e);
        }
    }

}